
#ifndef TRACE_GEO_ROUTE_H
#define TRACE_GEO_ROUTE_H

#include <cassert>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <linux/ip.h>
#include <linux/icmp.h>
#include <unistd.h>
#include <iostream>
#include <string>
#include <cstring>
#include <net/if.h>
#include <sys/ioctl.h>
#include "Utility.h"
#include <fstream>
#include <vector>
#include <iterator>
#include <algorithm>
#include <map>
#include <set>

#define MAX_TTL 64
#define ECHO_TIMEOUT 3
#define ECHO_INTERVAL 1


typedef std::map<std::string, std::string> StringMapType;
typedef std::set<std::string> StringSetType;


int GeoFromName(std::vector<std::pair<std::string, std::string> >& hostnames);


void ToLower(char& c)
{
	c = tolower(c);
}


int TraceGeoRoute(std::string dstAddr)
{

    struct iphdr *ip, *ip_reply;
    struct icmphdr *icmp, *icmp_reply;
    struct sockaddr_in connection;
    char* packet;
    char* buffer;
    int sockfd;
    
	std::string srcAddr;
	GetSourceIP(srcAddr);

	std::cout << "Source address: " << srcAddr << std::endl;
	std::cout << "Destination address: " << dstAddr << std::endl << std::endl;
    
  	packet = new char[sizeof(struct iphdr) + sizeof(struct icmphdr)];
	buffer = new char[sizeof(struct iphdr) + sizeof(struct icmphdr)];

    ip = (struct iphdr*) packet;
    icmp = (struct icmphdr*) (packet + sizeof(struct iphdr));


	//
	// The ICMP packet
	//
    icmp->type			= ICMP_ECHO;
    icmp->code			= 0; // error message code, none = 0
    icmp->un.echo.id		= 0;
    icmp->un.echo.sequence	= 0; // sequence # if more than one are sent
    icmp-> checksum		= in_cksum((unsigned short *)icmp, sizeof(struct icmphdr));

	//
	// The IP header
	//
    ip->ihl			= 5; // header length = 5*4 = 20 bytes
    ip->version			= 4; // IPv4
    ip->tos			= 0; // normal service
    ip->tot_len			= sizeof(struct iphdr) + sizeof(struct icmphdr);
    ip->id			= htons(random()); // single datagram, ID doesn't matter
    //ip->ttl			= 255; // max is 255
    ip->protocol		= IPPROTO_ICMP;
    ip->saddr			= inet_addr(srcAddr.c_str());
    ip->daddr			= inet_addr(dstAddr.c_str());
    //ip->check			= in_cksum((unsigned short *)ip, sizeof(struct iphdr));
 
    

   
    connection.sin_family = AF_INET;
    connection.sin_addr.s_addr = inet_addr(dstAddr.c_str());
	socklen_t addrlen;

	int selectVal, whoisVal, nslookupVal;
	unsigned int echoCounter, timeoutCounter = 0;
	std::vector<std::pair<std::string, std::string> > hostnames;

	for(echoCounter = 0; echoCounter < MAX_TTL; ++echoCounter)	
	{
		sleep(ECHO_INTERVAL);

		ip->ttl = echoCounter+1;
		ip->check = in_cksum((unsigned short*)ip, sizeof(struct iphdr));


		//
		// create raw socket each iteration
		// otherwise the server doesn't seem to respond
		// something to do with the port, echo.id and echo.seq
		//
		if ((sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) == -1)
		{
			perror("socket");
			exit(EXIT_FAILURE);
		}
		
		// notify the lower layers of the custom headers    
		int optval = 1;
		if(setsockopt(sockfd, IPPROTO_IP, IP_HDRINCL, &optval, sizeof(int)) == -1)
		{
			std::cout << "Unable to set the IP_HDRINCL option. Exiting..." << std::endl;
			close(sockfd);
			return -1;
		}


  		// send packet   
		int snd = sendto(sockfd, packet, ip->tot_len, 0, (struct sockaddr *)&connection, sizeof(struct sockaddr));
		if(snd == -1)
		{
			perror("sendto");
			close(sockfd);
			return -1;
		}
		/*else
			std::cout << "Sent " << snd << " bytes of attempt# " << echoCounter << " to " << dstAddr << std::endl;
		*/
		fd_set fds;
		FD_ZERO(&fds);
		FD_SET(sockfd, &fds);
   		struct timeval tv;
		tv.tv_sec = ECHO_TIMEOUT;
		tv.tv_usec = 0;
		selectVal = select(sockfd+1, &fds, NULL, NULL, &tv);
		if(selectVal == -1)
		{
			perror("select");
			close(sockfd);
			return -1;
		}
		else if(selectVal == 0)
		{
			std::cout << "Timeout occured for attempt# " << echoCounter << std::endl;
			close(sockfd);

			++timeoutCounter;
			if(timeoutCounter == 5) // stop after a while
				break;

			continue;
		}

		// select did not timeout this time
		timeoutCounter = 0;
		
		addrlen = sizeof(connection);
		int recv = recvfrom(sockfd, buffer, sizeof(struct iphdr) + sizeof(struct icmphdr), 0, (struct sockaddr *)&connection, &addrlen);
		if (recv == -1)
		{
			perror("recv");
		}
		else if (recv == 0)
		{
			std::cout << "Remote side closed the socket!" << std::endl;
			return -1;
		}
		else
		{
			std::string replyerIP(inet_ntoa(((struct sockaddr_in*)&connection)->sin_addr));
			//std::cout << "Received " << recv << " byte reply from " << replyerIP.c_str() << std::endl;
			ip_reply = (struct iphdr*) buffer;
			icmp_reply = (struct icmphdr*) (buffer + sizeof(struct iphdr));
			//std::cout << "ID: " << ntohs(ip_reply->id) << std::endl;
			//std::cout << "TTL: " << (int)ip_reply->ttl << std::endl;
			//std::cout << "ICMP type: " << (int)icmp_reply->type << std::endl;

			// store the IP address for future printing
			hostnames.push_back(std::make_pair(replyerIP, std::string()));
			
			// now do an nslookup on the replyerIP for geo hints
			char nslookupCmd[100];
			sprintf(nslookupCmd, "nslookup %s > nslookupop%02d", replyerIP.c_str(), (int)hostnames.size()-1);
			nslookupVal = system(nslookupCmd);

			if((int)icmp_reply->type == ICMP_ECHOREPLY)
				break;
		}
    	close(sockfd);
	}
	
	
	// echoCounter now contains the number of whois output files
	char fileName[20];
	char line0[500];
	char line1[500];
	int previousAS = 0, newAS = 0;
	size_t pos;
	const char SEPARATOR = '|';
	char rmCmd[50];

	for(unsigned int nslookupCtr = 0; nslookupCtr < hostnames.size(); ++nslookupCtr)
	{
		sprintf(fileName, "nslookupop%02d", nslookupCtr);
		std::ifstream nslookupFile(fileName);
		
		std::vector<std::string> tokens;
		std::string token;
		while(1)	
		{
			nslookupFile >> token;
			if(nslookupFile.good() == false)
				break;

			tokens.push_back(token);
		}

		for(unsigned int i = 0; i < tokens.size(); ++i)
		{
			if(tokens[i].compare("name") == 0)
				if(tokens[i+1].compare("=") == 0)
				{
					//hostnames.push_back(tokens[i+2]);
					hostnames[nslookupCtr].second = tokens[i+2];
					break;
				}
		}
		nslookupFile.close();

		// comment out the below two lines if you need the nslookup output files
		sprintf(rmCmd, "rm %s", fileName);
		system(rmCmd);

	}

	
	for(unsigned int i = 0; i < hostnames.size(); ++i)
	{
		std::for_each(hostnames[i].second.begin(), hostnames[i].second.end(), &ToLower);
	}

	//std::copy(hostnames.begin(), hostnames.end(), std::ostream_iterator<std::string>(std::cout, "\n"));
	GeoFromName(hostnames);

	// delete the dynamically allocated memory
	delete[] packet;
	delete[] buffer;

    return 0;
}


int GeoFromName(std::vector<std::pair<std::string, std::string> >& hostnames)
{
	char line[100];
	std::ifstream countryCodesFile("country-codes.txt");
	StringMapType countryCodeMap;
	StringSetType countryNameSet;
	std::string lineStr;
	int pos;
	while(1)
	{
		countryCodesFile.getline(line, 100);
		if(countryCodesFile.good() == false)
			break;
		lineStr.assign(line);
		std::for_each(lineStr.begin(), lineStr.end(), &ToLower);
		pos = lineStr.find_first_of(' ');
		countryCodeMap.insert(std::make_pair(lineStr.substr(0, pos), lineStr.substr(pos+1)));
		countryNameSet.insert(lineStr.substr(pos+1));
	}
	countryCodesFile.close();

	std::ifstream stateCodesFile("usps_state_abbrev");
	StringMapType stateCodeMap;
	StringSetType stateNameSet;
	while(1)
	{
		stateCodesFile.getline(line, 100);
		if(stateCodesFile.good() == false)
			break;
		lineStr.assign(line);
		std::for_each(lineStr.begin(), lineStr.end(), &ToLower);
		pos = lineStr.find_last_of(' ');
		stateCodeMap.insert(std::make_pair(lineStr.substr(pos+1), lineStr.substr(0, pos)));
		stateNameSet.insert(lineStr.substr(0, pos));
	}
	stateCodesFile.close();

#if 0
	std::cout << "Size of country codes: " << countryCodeMap.size() << std::endl;
	std::cout << "Size of state codes: " << stateCodeMap.size() << std::endl;

	for(StringMapType::const_iterator tempIter = countryCodeMap.begin();
			tempIter != countryCodeMap.end(); ++tempIter)
	{
		std::cout << tempIter->first << " " << tempIter->second << std::endl;
	}

	for(StringMapType::const_iterator tempIter = stateCodeMap.begin();
			tempIter != stateCodeMap.end(); ++tempIter)
	{
		std::cout << tempIter->first << " " << tempIter->second << std::endl;
	}
#endif

	std::ifstream cityNamesFile("major_cities_of_the_world.txt");
	StringSetType cityNameSet;
	while(1)
	{
		cityNamesFile.getline(line, 100);
		if(cityNamesFile.good() == false)
			break;
		lineStr.assign(line);
		std::for_each(lineStr.begin(), lineStr.end(), &ToLower);
		cityNameSet.insert(lineStr);
	}
	cityNamesFile.close();


	// assume all airport codes are 3 letter
	std::ifstream airportCodesFile("airport_codes");
	StringMapType airportCodeMap;
	while(1)
	{
		airportCodesFile.getline(line, 100);
		if(airportCodesFile.good() == false)
			break;
		lineStr.assign(line);
		std::for_each(lineStr.begin(), lineStr.end(), &ToLower);

		pos = lineStr.find_first_of('(');
		assert(pos != std::string::npos);

		airportCodeMap.insert(std::make_pair(lineStr.substr(pos+1, 3), lineStr.substr(0, pos)));
	}
	airportCodesFile.close();

	std::ifstream sprintCityCodesFile("sprintlink_city_codes");
	StringMapType sprintCityCodeMap;
	while(1)
	{
		sprintCityCodesFile.getline(line, 100);
		if(sprintCityCodesFile.good() == false)
			break;
		lineStr.assign(line);
		pos = lineStr.find_first_of(' ');
		sprintCityCodeMap.insert(std::make_pair(lineStr.substr(0, pos), lineStr.substr(pos+1)));
	}
	sprintCityCodesFile.close();


	///////////////////////
	// ALL DATASETS LOADED
	//
	// NOW BEGIN ALGORITHM
	///////////////////////


	StringMapType::const_iterator tempMapIter;
	StringSetType::const_iterator tempSetIter;
	bool geoFound;
	for(unsigned int i = 0; i < hostnames.size(); ++i)
	{
		std::string& hostname = hostnames[i].second;
		std::cout << hostnames[i].first << " = " << hostnames[i].second << " ==> ";

		geoFound = false;

		//
		// find city code if its a sprintlink router
		// sl-[ bb | gw | dr | st | pe | crs ]##-xxx.sprintlink.net
		//
		if((pos = hostname.find("sprintlink")) != std::string::npos)
		{
			for(tempMapIter = sprintCityCodeMap.begin(); tempMapIter != sprintCityCodeMap.end(); ++tempMapIter)
			{
				if((pos = hostname.find(tempMapIter->first)) != std::string::npos)
				{
					// lets avoid false positives here as well using the same assumption / rule
					if((pos == 0) || (hostname[pos-1] < 'a') || (hostname[pos-1] > 'z'))
					{
						pos += tempMapIter->first.length(); // after the code
						if((pos == hostname.length()) || (hostname[pos] < 'a') || (hostname[pos] > 'z'))
						{
							geoFound = true;
							std::cout << tempMapIter->second << std::endl;
							break;
						}
					}
				}
			}
		}

		if(geoFound)
			continue;

		// find airport code, assume all airport codes are 3 letter
		for(tempMapIter = airportCodeMap.begin(); tempMapIter != airportCodeMap.end(); ++tempMapIter)
		{
			if((pos = hostname.find(tempMapIter->first)) != std::string::npos)
			{
				// this 3 letter airport code could be a false positive,
				// as a general rule, assume that there can't be an alphabet
				// before and after the code
				if((pos == 0) || (hostname[pos-1] < 'a') || (hostname[pos-1] > 'z'))
				{
					pos += 3; // after the code
					if((pos == hostname.length()) || (hostname[pos] < 'a') || (hostname[pos] > 'z'))
					{
						geoFound = true;
						std::cout << tempMapIter->second << std::endl;
						break;
					}
				}
			}
		}

		if(geoFound)
			continue;

		// find city name
		for(tempSetIter = cityNameSet.begin(); tempSetIter != cityNameSet.end(); ++tempSetIter)
		{
			if(hostname.find(*tempSetIter) != std::string::npos)
			{
				geoFound = true;
				std::cout << *tempSetIter << std::endl;
				break;
			}
		}

		if(geoFound)
			continue;

		// find the extension between the last two periods, thats the country code
		pos = hostname.find_last_of('.', hostname.length()-2); // we want to ignore the last period
		if((pos != std::string::npos) && (hostname.length() - 2 - pos == 2))
		{
			// the root domain is some country code
			std::string rootDomain(hostname.substr(pos+1, 2));
			if((tempMapIter = countryCodeMap.find(rootDomain)) != countryCodeMap.end())
			{
				geoFound = true;
				std::cout << tempMapIter->second << std::endl;
				continue;
			}
		}
		
		if(geoFound)
			continue;

		// find state name
		for(tempSetIter = stateNameSet.begin(); tempSetIter != stateNameSet.end(); ++tempSetIter)
		{
			if(hostname.find(*tempSetIter) != std::string::npos)
			{
				geoFound = true;
				std::cout << *tempSetIter << std::endl;
				break;
			}
		}

		if(geoFound)
			continue;

		// find country name
		for(tempSetIter = countryNameSet.begin(); tempSetIter != countryNameSet.end(); ++tempSetIter)
		{
			if(hostname.find(*tempSetIter) != std::string::npos)
			{
				geoFound = true;
				std::cout << *tempSetIter << std::endl;
				break;
			}
		}

		if(geoFound)
			continue;

		std::cout << std::endl; // no location found :(
	}

	return 0;
}

#endif // TRACE_GEO_ROUTE_H






















